import pennylane as qml
import torch
import torch.nn as nn
import torch.optim as optim

# TODO: 外部定义n_qubits
n_qubits = 4
dev = qml.device("default.qubit", wires=n_qubits)	  # w/o noise
# dev = qml.device('qiskit.aer', wires=n_qubits, noise_model=noise_model)	  # w noise

@qml.qnode(dev)
def qnode(inputs, w00, w01, w02, w03,
		  x00, x01, x02, x03,
		  w10, w11, w12, w13,
		  x10, x11, x12, x13,):
	# print(inputs.shape)
	# for rx in range(n_qubits):
	# 	print(inputs[rx].shape)
	# 	qml.RX(inputs[rx], wires=rx)
	# print('=====')
	qml.AngleEmbedding(inputs, wires=range(n_qubits))

	qml.Rot(*w00, wires=0)
	qml.Rot(*w01, wires=1)
	qml.Rot(*w02, wires=2)
	qml.Rot(*w03, wires=3)

	qml.CRX(x00, wires=[0,1])
	qml.CRX(x01, wires=[1,2])
	qml.CRX(x02, wires=[2,3])
	qml.CRX(x03, wires=[3,0])

	# for rx in range(n_qubits):
	# 	qml.RX(inputs[rx], wires=rx)
	
	qml.AngleEmbedding(inputs, wires=range(n_qubits))

	qml.Rot(*w10, wires=0)
	qml.Rot(*w11, wires=1)
	qml.Rot(*w12, wires=2)
	qml.Rot(*w13, wires=3)

	qml.CRX(x10, wires=[0,1])
	qml.CRX(x11, wires=[1,2])
	qml.CRX(x12, wires=[2,3])
	qml.CRX(x13, wires=[3,0])

	return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

weight_shapes = {"w00": 3, "w01": 3, "w02": 3, "w03": 3, 
				 "x00": 1, "x01": 1, "x02": 1, "x03": 1, 
				 "w10": 3, "w11": 3, "w12": 3, "w13": 3, 
				 "x10": 1, "x11": 1, "x12": 1, "x13": 1, }

class QuantumLayer(nn.Module):
	def __init__(self):
		super().__init__()
		# VQC layer
		self.qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)

	def forward(self, x):
		# b = x.shape[0] # batch size
		# print(f'=== forward: {x.shape}')
		# x = x.view(b, 4)
		# print(f'=== forward: {x.shape}')
		x = self.qlayer(x)
		# out = F.log_softmax(out, dim=1)
		return x.to(torch.float32)

if __name__=="__main__":
	# 创建Sequential模型并嵌套自定义模块
	model = nn.Sequential(
		QuantumLayer(),
		nn.Linear(4, 9)
	)

	# 定义损失函数和优化器
	criterion = nn.MSELoss()
	optimizer = optim.Adam(model.parameters(), lr=0.1)

	# 进行训练
	input_data = torch.randn(2, 4)
	target = torch.randint(0, 10, (2, 9)).float()

	print(input_data)
	print(target)
	for epoch in range(1):
		# 前向传播
		output = model(input_data)
		print(output)
		loss = criterion(output, target)

		# 反向传播和参数更新
		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

		# 打印训练信息
		print(f"Epoch: {epoch+1}, Loss: {loss.item()}")
